import cv2 as cv
from ultralytics import YOLO
import math
import os
import time

# 姿态估计和手机检测模型路径
pose_model_path = 'yolov8n-pose.pt'
phone_model_path = 'yolov8n.pt'

# 加载模型
pose_model = YOLO(pose_model_path)
phone_model = YOLO(phone_model_path)

# 输入图片文件夹和输出保存文件夹
input_folder = "D:/python/v8_mine/dataset/"
output_folder = "D:/python/v8_mine/output_images/"

# 确保输出文件夹存在
if not os.path.exists(output_folder):
    os.makedirs(output_folder)

# 图片文件扩展名
image_extensions = image_extensions = ('.jpg', '.png', '.jpeg')

# 遍历输入文件夹中的图片
for filename in os.listdir(input_folder):
    if filename.lower().endswith(image_extensions):
        # 读取图片
        filepath = os.path.join(input_folder, filename)
        image = cv.imread(filepath)

        # 检测手机
        phone_result = phone_model(image)
        phone_result = phone_result[0]
        phone_boxes = phone_result.boxes
        phone_detected = False
        class_ids = phone_boxes.cls.cpu().numpy().astype(int)
        if 67 in class_ids:  # 检测到手机
            for phone_box in phone_boxes:
                xyxy_tensor = phone_box.xyxy[0]  # 提取cellphone边界框的xyxy信息
                xyxy_coords = xyxy_tensor.cpu().numpy()
                x1, y1, x2, y2 = xyxy_coords
            phone_detected = True

            # 检测姿态
            if phone_detected:
                pose_results = pose_model(image)
                results = pose_model(image)
                keypoints = results[0].keypoints

                xy_cpu = keypoints.xy.cpu() if keypoints.xy.device != 'cpu' else keypoints.xy

                # 根据官方模型的说明书，从17个点位中提取左右臂共6个点位
                shoulder_x_r, shoulder_y_r = xy_cpu[0, 6, 0].item(), xy_cpu[0, 6, 1].item()
                elbow_x_r, elbow_y_r = xy_cpu[0, 8, 0].item(), xy_cpu[0, 8, 1].item()
                wrist_x_r, wrist_y_r = xy_cpu[0, 10, 0].item(), xy_cpu[0, 10, 1].item()
                shoulder_x_l, shoulder_y_l = xy_cpu[0, 5, 0].item(), xy_cpu[0, 5, 1].item()
                elbow_x_l, elbow_y_l = xy_cpu[0, 7, 0].item(), xy_cpu[0, 7, 1].item()
                wrist_x_l, wrist_y_l = xy_cpu[0, 9, 0].item(), xy_cpu[0, 9, 1].item()

                # 根据官方模型，从17个点位提取髋关节位置，通过判断肩至髋的长度与肩宽的关系，判断是否处于侧身
                hip_y_l = xy_cpu[0, 11, 1]
                hip_y_r = xy_cpu[0, 12, 1]

                # 计算左右臂的四个向量
                vector_A_x_r = elbow_x_r - shoulder_x_r
                vector_A_y_r = elbow_y_r - shoulder_y_r

                vector_B_x_r = wrist_x_r - elbow_x_r
                vector_B_y_r = wrist_y_r - elbow_y_r

                vector_A_x_l = elbow_x_l - shoulder_x_l
                vector_A_y_l = elbow_y_l - shoulder_y_l

                vector_B_x_l = wrist_x_l - elbow_x_l
                vector_B_y_l = wrist_y_l - elbow_y_l

                # 计算左右臂的大臂、小臂向量乘积
                dot_product_r = vector_A_x_r * vector_B_x_r + vector_A_y_r * vector_B_y_r
                dot_product_l = vector_A_x_l * vector_B_x_l + vector_A_y_l * vector_B_y_l

                # 计算长度
                magnitude_A_r = math.sqrt(vector_A_x_r ** 2 + vector_A_y_r ** 2)
                magnitude_B_r = math.sqrt(vector_B_x_r ** 2 + vector_B_y_r ** 2)
                magnitude_A_l = math.sqrt(vector_A_x_l ** 2 + vector_A_y_l ** 2)
                magnitude_B_l = math.sqrt(vector_B_x_l ** 2 + vector_B_y_l ** 2)

                # 保障计算时分母不为0
                if magnitude_A_r * magnitude_B_r == 0:
                    continue
                if magnitude_A_l * magnitude_B_l == 0:
                    continue

                # 计算cos值
                cosine_angle_r = dot_product_r / (magnitude_A_r * magnitude_B_r)
                cosine_angle_l = dot_product_l / (magnitude_A_l * magnitude_B_l)

                # 计算角度
                angle_in_degrees_r = math.degrees(math.acos(cosine_angle_r))
                angle_in_degrees_l = math.degrees(math.acos(cosine_angle_l))

                # 找出最接近九十度的夹角
                closest_angle = min(angle_in_degrees_r, angle_in_degrees_l, key=lambda x: abs(x - 98.12))

                # 首先计算手机的中心位置
                x_phone_mid = (x1 + x2) / 2
                y_phone_mid = (y1 + y2) / 2

                # magnitude_B_l和magnitude_B_r为左右侧的小臂长度
                # 计算手机中心到两手腕的距离，取最小的一个作为所需
                phone_to_leftwrist = math.sqrt((x_phone_mid - wrist_x_l) ** 2 + (y_phone_mid - wrist_y_l) ** 2)
                phone_to_rightwrist = math.sqrt((x_phone_mid - wrist_x_r) ** 2 + (y_phone_mid - wrist_y_r) ** 2)
                phone_to_wirst = min(phone_to_rightwrist, phone_to_leftwrist)

                # 判断是否为侧身：肩宽/髋宽
                height_l = shoulder_y_l - hip_y_l
                height_r = shoulder_y_r - hip_y_r
                delta_y_hiptoshoulder = max(height_l, height_r)
                delta_x_shoulder = math.sqrt((shoulder_x_r - shoulder_x_l) ** 2 + (shoulder_y_r - shoulder_y_l) ** 2)
                bili = delta_x_shoulder / delta_y_hiptoshoulder

                # 检查角度是否在指定范围内，并且手机被检测到，且是该人在走路玩手机
                if 51.71 <= closest_angle <= 144.53 and phone_to_wirst <= magnitude_B_l:
                    # 因为存不到照片，减弱判断范围
                    if -0.36139454 <= abs(bili) <= 0.17988266:  # 不加有17/296，加了有3/296
                        print("检测到走路玩手机的同学，已经保存")
                        output_filepath = os.path.join(output_folder, 'result_' + filename)
                        cv.imwrite(output_filepath, image)

print("处理完成。")
